#include <pinocchio/fwd.hpp>
#include <ocs2_oc/oc_problem/OptimalControlProblem.h>
#include <ocs2_oc/rollout/TimeTriggeredRollout.h>
#include <ocs2_pinocchio_interface/PinocchioInterface.h>
#include <pinocchio/multibody/data.hpp>
#include <pinocchio/multibody/model.hpp>
#include <pinocchio/parsers/urdf.hpp>
#include <ocs2_core/cost/QuadraticStateInputCost.h>
#include <ocs2_core/cost/QuadraticStateCost.h>
#include <ocs2_core/constraint/LinearStateInputConstraint.h>
#include <ocs2_core/augmented_lagrangian/AugmentedLagrangian.h>
#include <ocs2_core/penalties/augmented/SlacknessSquaredHingePenalty.h>
#include <ocs2_core/initialization/DefaultInitializer.h>
#include <ocs2_core/thread_support/ExecuteAndSleep.h>
#include <ocs2_core/thread_support/SetThreadPriority.h>
#include <ocs2_core/misc/LoadData.h>
#include <ocs2_ddp/GaussNewtonDDP_MPC.h>
#include <ocs2_ros_interfaces/mpc/MPC_ROS_Interface.h>
#include <ocs2_mpc/MPC_MRT_Interface.h>

#include <ocs2_msgs/mpc_observation.h>
#include <ocs2_msgs/mpc_target_trajectories.h>
#include <ocs2_msgs/mpc_input.h>

#include <ocs2_ros_interfaces/common/RosMsgConversions.h>
#include <string>
#include <iostream>

#include <ros/ros.h>
#include <ros/package.h>

#include "cartpole_dynamic.h"

using namespace ocs2;
using namespace cartpole;

namespace cartpole {
    const size_t STATE_DIM = 8;
    const size_t INPUT_DIM = 1;
    bool init_observation = false;
}

int main(int argc, char **argv) {
    ros::init(argc, argv, "cartpole");
    ros::NodeHandle node_handle;

    std::string package_path = ros::package::getPath("cartpole_gazebo_ros");
    std::string cartpole_urdf = package_path + std::string("/model/cartpole_pinocchio.urdf");
    std::string config_file = package_path + std::string("/config/config.info");
    std::string lib_path = package_path + std::string("/auto_generator");

    // 模型 使用 pinocchio
    pinocchio::Model cartpole_model;
    pinocchio::urdf::buildModel(cartpole_urdf, cartpole_model);
    PinocchioInterface pinocchioInterface(cartpole_model);
    // Optimal control problem
    OptimalControlProblem OCP;
    //   1. dynamic
    CartpoleDynamic dynamic(pinocchioInterface);
    OCP.dynamicsPtr.reset(&dynamic);

    //   2. cost
    // 创建Q R矩阵
    matrix_t Q(STATE_DIM, STATE_DIM);
    matrix_t Qf(STATE_DIM, STATE_DIM);
    vector_t R(INPUT_DIM);
    loadData::loadEigenMatrix(config_file, "Q", Q);
    loadData::loadEigenMatrix(config_file, "R", R);
    loadData::loadEigenMatrix(config_file, "Q_final", Qf);

    OCP.costPtr->add("cartpole_cost", std::unique_ptr<StateInputCost>(new QuadraticStateInputCost(Q, R)));
    OCP.finalCostPtr->add("cartpole_final_cost", std::unique_ptr<StateCost>(new QuadraticStateCost(Qf)));
    //   3. constraint
    // Constraints
    // todo 查看不同 的约束算法
    auto getPenalty = [&]() {
        using penalty_type = augmented::SlacknessSquaredHingePenalty;
        penalty_type::Config boundsConfig(0.1, 1);
        return penalty_type::create(boundsConfig);
    };

    scalar_t max_input;
    loadData::loadCppDataType(config_file, "max_input", max_input);
    auto getConstraint = [&]() {
        constexpr size_t numIneqConstraint = 2;
        const vector_t e = (vector_t(numIneqConstraint) << max_input, max_input).finished();
        const vector_t D = (vector_t(numIneqConstraint) << 1.0, -1.0).finished();
        const matrix_t C = matrix_t::Zero(numIneqConstraint, STATE_DIM);
        return std::unique_ptr<StateInputConstraint>(new LinearStateInputConstraint(e, C, D));
    };
    OCP.inequalityLagrangianPtr->add("InputLimits", create(getConstraint(), getPenalty()));

    // 要有一个MPC
    rollout::Settings rolloutSetting;
    rolloutSetting = rollout::loadSettings(config_file, "rollout");
    TimeTriggeredRollout rollout(dynamic, rolloutSetting);
    mpc::Settings mpc_setting = mpc::loadSettings(config_file, "mpc");
    ddp::Settings ddp_setting = ddp::loadSettings(config_file, "ddp");
    DefaultInitializer initializer(INPUT_DIM);
    GaussNewtonDDP_MPC mpc(mpc_setting, ddp_setting, rollout, OCP, initializer);

    // 使用MPC MRT 进行控制 提高 控制 频率
    MPC_MRT_Interface mrt(mpc);
    mrt.initRollout(&rollout);
    SystemObservation systemObservation;
    TargetTrajectories initTarget{{0},
                                  {vector_t(STATE_DIM).setZero()},
                                  {vector_t(INPUT_DIM).setZero()}};

    auto rate = ros::Rate(mpc_setting.mrtDesiredFrequency_);
    // 创建两个subscriber
    // 必须指定一个变量接受，否则会被直接销毁，无法进入回调函数
    auto target_sub = node_handle.subscribe<ocs2_msgs::mpc_target_trajectories>(
            "/cartpole_target",
            1,
            [&mrt](const ocs2_msgs::mpc_target_trajectories::ConstPtr &msg) {
                auto target = ros_msg_conversions::readTargetTrajectoriesMsg(*msg);
                mrt.getReferenceManager().setTargetTrajectories(target);
            });
    auto observation_sub = node_handle.subscribe<ocs2_msgs::mpc_observation>(
            "/cartpole_mpc_observation",
            1,
            [&](const ocs2_msgs::mpc_observation::ConstPtr &msg) {
                auto observation = ros_msg_conversions::readObservationMsg(*msg);
                mrt.setCurrentObservation(observation);
                systemObservation = observation;
                if (!init_observation) {
                    init_observation = true;
                }
            });

    auto policyPub = node_handle.advertise<ocs2_msgs::mpc_input>("/cartpole_mpc_input", 1, true);

    mrt.resetMpcNode(initTarget);
    ROS_INFO_STREAM("Waiting for the initial observation ...");
    while (!init_observation) {
        ros::spinOnce();
        ros::WallRate(mpc_setting.mrtDesiredFrequency_).sleep();
    }
    mrt.advanceMpc();
    std::thread mpc_thread([&]() {
        executeAtRate([&]() {
            mrt.advanceMpc();
        }, []() {
            return true;
        }, mpc_setting.mpcDesiredFrequency_);
    });
    setThreadPriority(ddp_setting.threadPriority_, mpc_thread);

    while (!ros::isShuttingDown()) {
        // 运行 mrt run
        vector_t mpc_state(STATE_DIM);
        vector_t mpc_input(INPUT_DIM);
        size_t mode;

        ros::spinOnce();
        mrt.updatePolicy();

        if (mrt.isRolloutSet()) {
            mrt.rolloutPolicy(systemObservation.time,
                              systemObservation.state,
                              1.0 / mpc_setting.mrtDesiredFrequency_,
                              mpc_state,
                              mpc_input,
                              mode);
        } else {
            mrt.evaluatePolicy(systemObservation.time + 1.0 / mpc_setting.mrtDesiredFrequency_,
                               systemObservation.state,
                               mpc_state,
                               mpc_input,
                               mode);
        }
        // publish 出去
        ocs2_msgs::mpc_input out;
        out.value.resize(INPUT_DIM);
        out.value[0] = (float) mpc_input[0];
        policyPub.publish(out);
        rate.sleep();
    }
}